package main

import (
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/spf13/cobra"
)

var (
	subnet      string
	portStr     string
	concurrency int
)

func main() {
	var rootCmd = &cobra.Command{
		Use:   "scannet",
		Short: "Network scanner CLI tool",
		Run:   runScan,
	}

	rootCmd.Flags().StringVarP(&subnet, "subnet", "s", "", "Subnet to scan (e.g., 192.168.1.0/24)")
	rootCmd.Flags().StringVarP(&portStr, "ports", "p", "80", "Ports to scan (e.g., 80,443,3000-3005)")
	rootCmd.Flags().IntVarP(&concurrency, "concurrency", "c", 100, "Number of concurrent scans")

	if err := rootCmd.Execute(); err != nil {
		fmt.Println(err)
	}
}

// 计算子网内的所有 IP 地址
func calSubnet(subnet string) ([]string, error) {
	_, ipNet, err := net.ParseCIDR(subnet)
	if err != nil {
		return nil, err
	}

	var ips []string
	ip := make(net.IP, len(ipNet.IP))
	copy(ip, ipNet.IP)
	for ipNet.Contains(ip) {
		ips = append(ips, ip.String())
		inc(ip)
	}
	return ips, nil
}

// 递增 IP 地址
func inc(ip net.IP) {
	for j := len(ip) - 1; j >= 0; j-- {
		ip[j]++
		if ip[j] != 0 {
			break
		}
	}
}

// 解析端口范围
func parseData(dataString string) []int {
	var ports []int
	if dataString == "" {
		return ports
	}
	ranges := strings.Split(dataString, ",")
	for _, item := range ranges {
		if strings.Contains(item, "-") {
			var start, end int
			fmt.Sscanf(item, "%d-%d", &start, &end)
			for i := start; i <= end; i++ {
				ports = append(ports, i)
			}
		} else {
			var port int
			fmt.Sscanf(item, "%d", &port)
			ports = append(ports, port)
		}
	}
	return ports
}

// 扫描单个端口
func scanPort(targetHost string, port int, timeout time.Duration) (string, error) {
	address := fmt.Sprintf("%s:%d", targetHost, port)
	conn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		return "", nil
	}
	defer conn.Close()
	return address, nil
}

// 执行端口扫描
func runScan(cmd *cobra.Command, args []string) {
	ips, err := calSubnet(subnet)
	if err != nil {
		fmt.Println("无效的子网:", err)
		return
	}

	portList := parseData(portStr)
	openPorts := make([]string, 0)
	var wg sync.WaitGroup
	sem := make(chan struct{}, concurrency)

	for _, ip := range ips {
		for _, port := range portList {
			wg.Add(1)
			sem <- struct{}{}
			go func(ip string, port int) {
				defer wg.Done()
				defer func() { <-sem }()
				if result, err := scanPort(ip, port, 1*time.Second); err == nil && result != "" {
					openPorts = append(openPorts, result)
					fmt.Printf("%s 开放\n", result)
				}
			}(ip, port)
		}
	}

	wg.Wait()
	if len(openPorts) == 0 {
		fmt.Println("未找到开放的端口")
	} else {
		fmt.Println("开放的端口：", strings.Join(openPorts, ", "))
	}
}
